#!/usr/bin/env python3
"""
飞书 (Lark) OCR Tkinter 图形界面

功能：
- 输入 app_id / app_secret 并获取 tenant_access_token
- 选择单张图片或图片文件夹进行识别
- 调用飞书 OCR 接口识别图片文字
- 显示原始 JSON 响应与提取的文字
- 批量处理并导出 CSV

依赖：
- Python 3.8+
- pip install requests pillow

运行：
python feishu_ocr_tkinter.py

注意：请确保你的飞书应用已开启 OCR 权限并已配置为租户访问。
"""

import base64
import io
import os
import sys
import time
import threading
import csv
import requests
from PIL import Image, ImageTk
import tkinter as tk
from tkinter import ttk
from tkinter import filedialog, messagebox, scrolledtext

# ----------------- 配置 -----------------
FEISHU_TOKEN_URL = "https://open.feishu.cn/open-apis/auth/v3/tenant_access_token/internal"
FEISHU_OCR_URL = "https://open.feishu.cn/open-apis/optical_char_recognition/v1/image/basic_recognize"
# ----------------------------------------

IMAGE_EXTS = ('.png', '.jpg', '.jpeg', '.bmp', '.webp')

class FeishuOCRApp(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("飞书 OCR - Tkinter 图形界面")
        self.geometry("900x650")
        self.resizable(True, True)

        self.token_info = {}
        self.selected_path = None  # 文件或文件夹
        self.batch_mode = False

        self._build_ui()

    def _build_ui(self):
        frm_top = ttk.Frame(self, padding=8)
        frm_top.pack(side=tk.TOP, fill=tk.X)

        # 凭证输入
        ttk.Label(frm_top, text="app_id：").grid(row=0, column=0, sticky=tk.W)
        self.entry_app_id = ttk.Entry(frm_top, width=40)
        self.entry_app_id.grid(row=0, column=1, sticky=tk.W, padx=4)

        ttk.Label(frm_top, text="app_secret：").grid(row=1, column=0, sticky=tk.W)
        self.entry_app_secret = ttk.Entry(frm_top, width=40, show='*')
        self.entry_app_secret.grid(row=1, column=1, sticky=tk.W, padx=4)

        self.btn_get_token = ttk.Button(frm_top, text="获取 Token", command=self.get_token_thread)
        self.btn_get_token.grid(row=0, column=2, rowspan=2, padx=8)

        ttk.Label(frm_top, text="token（预览）：").grid(row=2, column=0, sticky=tk.W, pady=(8,0))
        self.entry_token = ttk.Entry(frm_top, width=80)
        self.entry_token.grid(row=2, column=1, columnspan=2, sticky=tk.W, padx=4, pady=(8,0))

        # 文件选择
        frm_files = ttk.Frame(self, padding=8)
        frm_files.pack(side=tk.TOP, fill=tk.X)

        self.btn_select_file = ttk.Button(frm_files, text="选择图片", command=self.select_file)
        self.btn_select_file.grid(row=0, column=0)
        self.btn_select_folder = ttk.Button(frm_files, text="选择文件夹（批量）", command=self.select_folder)
        self.btn_select_folder.grid(row=0, column=1, padx=6)
        self.lbl_selected = ttk.Label(frm_files, text="未选择任何文件", width=80)
        self.lbl_selected.grid(row=0, column=2, sticky=tk.W, padx=6)

        # 图片预览 + 控件
        frm_middle = ttk.Frame(self, padding=8)
        frm_middle.pack(side=tk.TOP, fill=tk.BOTH, expand=True)

        # 左：图片预览
        frm_left = ttk.LabelFrame(frm_middle, text="图片预览", padding=6)
        frm_left.pack(side=tk.LEFT, fill=tk.BOTH, expand=False)
        self.canvas = tk.Canvas(frm_left, width=360, height=360, bg="#f0f0f0")
        self.canvas.pack()

        # 右：结果显示与操作
        frm_right = ttk.Frame(frm_middle, padding=6)
        frm_right.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        self.btn_recognize = ttk.Button(frm_right, text="识别", command=self.recognize_thread)
        self.btn_recognize.grid(row=0, column=0, sticky=tk.W)

        self.btn_show_json = ttk.Button(frm_right, text="查看上次 JSON", command=self.show_last_json)
        self.btn_show_json.grid(row=0, column=1, sticky=tk.W, padx=6)

        self.btn_save_csv = ttk.Button(frm_right, text="保存批量 CSV", command=self.save_csv)
        self.btn_save_csv.grid(row=0, column=2, sticky=tk.W, padx=6)

        ttk.Label(frm_right, text="识别到的文字：").grid(row=1, column=0, columnspan=3, sticky=tk.W, pady=(8,0))
        self.txt_result = scrolledtext.ScrolledText(frm_right, height=16)
        self.txt_result.grid(row=2, column=0, columnspan=3, sticky=tk.NSEW)

        ttk.Label(frm_right, text="日志：").grid(row=3, column=0, columnspan=3, sticky=tk.W, pady=(8,0))
        self.txt_log = scrolledtext.ScrolledText(frm_right, height=8)
        self.txt_log.grid(row=4, column=0, columnspan=3, sticky=tk.NSEW)

        # 让右侧布局可伸缩
        frm_right.rowconfigure(2, weight=1)
        frm_right.rowconfigure(4, weight=0)
        frm_right.columnconfigure(2, weight=1)

        # 状态栏
        self.status_var = tk.StringVar(value="就绪")
        status = ttk.Label(self, textvariable=self.status_var, relief=tk.SUNKEN, anchor=tk.W)
        status.pack(side=tk.BOTTOM, fill=tk.X)

        # 内部存储：上次响应与批量结果
        self.last_json = None
        self.batch_results = []
        self.current_image_tk = None

    # ----------------- 工具函数 -----------------
    def log(self, *args):
        t = time.strftime('%Y-%m-%d %H:%M:%S')
        self.txt_log.insert(tk.END, f"[{t}] " + ' '.join(map(str, args)) + '')
        self.txt_log.see(tk.END)

    def set_status(self, text):
        self.status_var.set(text)

    # ----------------- 获取 Token -----------------
    def get_token_thread(self):
        threading.Thread(target=self.get_token, daemon=True).start()

    def get_token(self):
        app_id = self.entry_app_id.get().strip()
        app_secret = self.entry_app_secret.get().strip()
        if not app_id or not app_secret:
            messagebox.showwarning("缺少凭证", "请填写 app_id 和 app_secret。")
            return
        self.set_status("正在请求 token...")
        self.log("正在请求 token...")
        try:
            resp = requests.post(FEISHU_TOKEN_URL, json={"app_id": app_id, "app_secret": app_secret}, timeout=8)
            data = resp.json()
            if resp.status_code == 200 and (data.get('code') == 0 or data.get('tenant_access_token')):
                # 不同文档 wrapper 略有差异，做兼容处理
                token = data.get('tenant_access_token') or data.get('data', {}).get('tenant_access_token')
                if token:
                    self.token_info = {'token': token, 'fetched_at': time.time(), 'raw': data}
                    self.entry_token.delete(0, tk.END)
                    self.entry_token.insert(0, token)
                    self.set_status("已获取 token")
                    self.log("已获取 token")
                else:
                    self.log('响应中未包含 token', data)
                    messagebox.showerror('Token 错误', f'响应中未包含 token: {data}')
                    self.set_status('Token 错误')
            else:
                self.log('Token 请求失败', resp.status_code, resp.text)
                messagebox.showerror('Token 错误', f'状态: {resp.status_code}响应: {resp.text}')
                self.set_status('Token 错误')
        except Exception as e:
            self.log('请求 token 异常', e)
            messagebox.showerror('网络错误', str(e))
            self.set_status('网络错误')

    # ----------------- 文件选择与预览 -----------------
    def select_file(self):
        path = filedialog.askopenfilename(filetypes=[('图片文件', '*.png;*.jpg;*.jpeg;*.bmp;*.webp'), ('所有文件','*.*')])
        if not path:
            return
        self.batch_mode = False
        self.selected_path = path
        self.lbl_selected.config(text=path)
        self._preview_image(path)
        self.log('选择文件', path)

    def select_folder(self):
        folder = filedialog.askdirectory()
        if not folder:
            return
        self.batch_mode = True
        self.selected_path = folder
        self.lbl_selected.config(text=folder)
        self._preview_image(None)
        self.log('选择文件夹', folder)

    def _preview_image(self, path):
        self.canvas.delete('all')
        self.current_image_tk = None
        if not path or not os.path.exists(path):
            return
        try:
            img = Image.open(path)
            img.thumbnail((360, 360))
            self.current_image_tk = ImageTk.PhotoImage(img)
            self.canvas.create_image(180, 180, image=self.current_image_tk)
        except Exception as e:
            self.log('预览错误', e)

    # ----------------- 识别 -----------------
    def recognize_thread(self):
        threading.Thread(target=self._recognize, daemon=True).start()

    def _recognize(self):
        token = self.entry_token.get().strip() or self.token_info.get('token')
        if not token:
            messagebox.showwarning('没有 token', '请先获取 tenant_access_token。')
            return
        if not self.selected_path:
            messagebox.showwarning('未选择文件', '请先选择图片文件或文件夹。')
            return

        headers = {'Authorization': f'Bearer {token}', 'Content-Type': 'application/json'}

        # 单张图片
        if not self.batch_mode:
            path = self.selected_path
            self.set_status('正在识别图片...')
            self.log('识别中', path)
            try:
                base64_str = self._image_to_base64(path)
                if not base64_str:
                    raise ValueError('图片编码失败')
                payload = {"image": base64_str}
                resp = requests.post(FEISHU_OCR_URL, json=payload, headers=headers, timeout=20)
                data = resp.json()
                self.last_json = data
                text = self._extract_text_from_response(data)
                self.txt_result.delete('1.0', tk.END)
                self.txt_result.insert(tk.END, text)
                self.log('识别结果', data.get('code'), data.get('msg'))
                self.set_status('完成')
            except Exception as e:
                self.log('识别异常', e)
                messagebox.showerror('识别错误', str(e))
                self.set_status('错误')

        # 批量文件夹
        else:
            folder = self.selected_path
            files = [os.path.join(folder, f) for f in os.listdir(folder) if f.lower().endswith(IMAGE_EXTS)]
            if not files:
                messagebox.showinfo('没有图片', '文件夹中未找到图片')
                return
            self.batch_results = []
            total = len(files)
            self.set_status(f'开始批量处理（{total} 张图片）')
            for idx, path in enumerate(files, 1):
                self.set_status(f'处理 {idx}/{total}: {os.path.basename(path)}')
                self.log('批量处理', idx, path)
                try:
                    base64_str = self._image_to_base64(path)
                    payload = {"image": base64_str}
                    resp = requests.post(FEISHU_OCR_URL, json=payload, headers=headers, timeout=25)
                    data = resp.json()
                    text = self._extract_text_from_response(data)
                    self.batch_results.append({'file': path, 'text': text, 'raw': data})
                    # 实时显示当前识别结果
                    self.txt_result.delete('1.0', tk.END)
                    self.txt_result.insert(tk.END, f"{os.path.basename(path)}:{text}")
                except Exception as e:
                    self.log('批量识别错误', e)
                    self.batch_results.append({'file': path, 'text': '', 'raw': {'error': str(e)}})
                time.sleep(0.2)  # 轻微延迟以避免触发速率限制
            self.set_status('批量处理完成')
            messagebox.showinfo('完成', f'已处理 {len(files)} 张图片。你可以保存为 CSV。')

    def _image_to_base64(self, path):
        try:
            with open(path, 'rb') as f:
                data = f.read()
            size_mb = len(data) / (1024 * 1024)
            if size_mb > 5:
                self.log('警告：图片大于 5MB，建议压缩或缩小分辨率', f'{size_mb:.2f}MB')
            b64 = base64.b64encode(data).decode('utf-8')
            return b64
        except Exception as e:
            self.log('图片编码错误', e)
            return None

    def _extract_text_from_response(self, data):
        # 飞书返回的结构可能不同，这里做兼容提取
        try:
            if not data:
                return ''
            if isinstance(data, dict):
                inner = data.get('data') or data
                text_list = inner.get('text_list') if isinstance(inner, dict) else None
                if text_list:
                    return ''.join(text_list)
                regions = inner.get('regions') if isinstance(inner, dict) else None
                if regions:
                    parts = []
                    for r in regions:
                        lines = r.get('lines') or []
                        for ln in lines:
                            txt = ln.get('text') or ln.get('words') or ''
                            parts.append(txt)
                    return ''.join(parts)
                return str(data)
            else:
                return str(data)
        except Exception as e:
            self.log('提取文本错误', e)
            return str(data)

    def show_last_json(self):
        if not self.last_json:
            messagebox.showinfo('无数据', '尚无上次响应记录')
            return
        top = tk.Toplevel(self)
        top.title('上次 JSON 响应')
        txt = scrolledtext.ScrolledText(top, width=100, height=40)
        txt.pack(fill=tk.BOTH, expand=True)
        import json
        txt.insert(tk.END, json.dumps(self.last_json, ensure_ascii=False, indent=2))

    def save_csv(self):
        if not self.batch_results:
            messagebox.showinfo('无结果', '当前没有批量识别结果。请先运行批量识别。')
            return
        csv_path = filedialog.asksaveasfilename(defaultextension='.csv', filetypes=[('CSV 文件','*.csv')])
        if not csv_path:
            return
        try:
            with open(csv_path, 'w', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                writer.writerow(['file', 'text'])
                for item in self.batch_results:
                    writer.writerow([item.get('file'), item.get('text')])
            messagebox.showinfo('已保存', f'CSV 已保存到 {csv_path}')
            self.log('CSV 已保存', csv_path)
        except Exception as e:
            self.log('CSV 保存错误', e)
            messagebox.showerror('保存错误', str(e))


if __name__ == '__main__':
    app = FeishuOCRApp()
    app.mainloop()
